#%%
import numpy as np, matplotlib.pyplot as plt, seaborn as sns, pandas as pd
%matplotlib inline

grains = ['gb', 'ng', 'oc', 'gg', 'ls']
grains_2022 = ['tg', 'ng2', 'ng3', 'ng4', 'ng5', 'sf']

df_density = pd.read_csv('main_granular_materials/1_var_density.csv').set_index('Unnamed: 0').rename_axis(None)
df_shape = pd.read_csv('main_granular_materials/2_var_shape.csv').set_index('Unnamed: 0').rename_axis(None)
df_ars = pd.read_csv('main_granular_materials/3_var_ARS.csv').set_index('Unnamed: 0').rename_axis(None)
df_vels = pd.read_csv('main_granular_materials/4_var_settling_velocities.csv').set_index('Unnamed: 0').rename_axis(None)
df_VES = pd.read_csv('main_granular_materials/5_var_VES_settling_velocity.csv').set_index('Unnamed: 0').rename_axis(None)

df_density_2022 = pd.read_csv('extended_methods_granular_materials/1_var_density_2022.csv').set_index('Unnamed: 0').rename_axis(None)
df_shape_2022 = pd.read_csv('extended_methods_granular_materials/2_var_shape_2022.csv').set_index('Unnamed: 0').rename_axis(None)
df_ars_2022 = pd.read_csv('extended_methods_granular_materials/3_var_ARS_2022.csv').set_index('Unnamed: 0').rename_axis(None)
df_vels_2022 = pd.read_csv('extended_methods_granular_materials/4_var_settling_velocities_2022.csv').set_index('Unnamed: 0').rename_axis(None)
df_VES_2022 = pd.read_csv('extended_methods_granular_materials/5_var_VES_settling_velocity_2022.csv').set_index('Unnamed: 0').rename_axis(None)

# %%
# calculate cstar and mustar

bed_angle = 3.5
aor_gb = 24
rof = 995
g = 9.81
cstar, dcstar, mustar, dmustar, ratio, dratio, Co, CDsettle, CD = {}, {}, {}, {}, {}, {}, {}, {}, {}
for grain in grains:
    cstar[grain] = ((df_VES.ws_sphere[grain])/(df_vels.vel_mean[grain]))**2 * df_shape.CSF[grain]
    dcstar[grain] = cstar[grain] * np.sqrt((2*df_vels.dvel_mean[grain]/df_vels.vel_mean[grain])**2 + (df_shape.dCSF[grain]/df_shape.CSF[grain])**2)

    mugb = (np.tan(np.radians(aor_gb)) - np.tan(np.radians(bed_angle)))
    mustar[grain] = (np.tan(np.radians(df_ars.mean_angle[grain])) - np.tan(np.radians(bed_angle)))/mugb
    dmustar[grain] = ((1 + np.tan(np.radians(df_ars.mean_angle[grain]))**2)/mugb) * np.radians(df_ars.mean_angle[grain]) * (df_ars.std_angle[grain]/df_ars.mean_angle[grain])
    ratio[grain] = cstar[grain]/mustar[grain]
    dratio[grain] = ratio[grain] * np.sqrt((dcstar[grain]/cstar[grain])**2 + (dmustar[grain]/mustar[grain])**2)
    
    Co[grain] = df_VES.CD_VES[grain]
    CDsettle[grain] = 4*(df_density.total_grain_ro[grain]/rof - 1)*g*df_shape.dn[grain] / (3*df_vels.vel_mean[grain]**2)
    CD[grain] = df_vels.CD[grain]
    # cstar[grain] = ((df_VES.ws_sphere[grain])/(df_vels.vel_mean[grain]))**2 * df_shape.CSF[grain]
    # dcstar[grain] = cstar[grain] * np.sqrt((2*df_VES.dws_sphere[grain]/df_VES.ws_sphere[grain])**2 + (2*df_vels.dvel_mean[grain]/df_vels.vel_mean[grain])**2 + (df_shape.dCSF[grain]/df_shape.CSF[grain])**2)
    # mugb = (np.tan(np.radians(aor_gb)) - np.tan(np.radians(bed_angle)))
    # mustar[grain] = (np.tan(np.radians(df_ars.mean_angle[grain])) - np.tan(np.radians(bed_angle)))/mugb
df_1 = pd.DataFrame({'Co': df_VES.CD_VES, 'CD': df_vels.CD, 'mustar': mustar, 'Cstar': cstar, 'dmustar': dmustar, 'dcstar': dcstar, 'ratio': ratio, 'dratio': dratio})

cstar, dcstar, mustar, dmustar, ratio, dratio, Co, CDsettle, CD = {}, {}, {}, {}, {}, {}, {}, {}, {}
for grain in grains_2022:
    cstar[grain] = ((df_VES_2022.ws_sphere[grain])/(df_vels_2022.vel_mean[grain]))**2 * df_shape_2022.CSF[grain]
    dcstar[grain] = cstar[grain] * np.sqrt((2*df_vels_2022.dvel_mean[grain]/df_vels_2022.vel_mean[grain])**2 + (df_shape_2022.dCSF[grain]/df_shape_2022.CSF[grain])**2)
    mugb = (np.tan(np.radians(aor_gb)) - np.tan(np.radians(bed_angle)))
    mustar[grain] = (np.tan(np.radians(df_ars_2022.mean_angle[grain])) - np.tan(np.radians(bed_angle)))/mugb
    dmustar[grain] = ((1 + np.tan(np.radians(df_ars_2022.mean_angle[grain]))**2)/mugb) * np.radians(df_ars_2022.mean_angle[grain]) * (df_ars_2022.std_angle[grain]/df_ars_2022.mean_angle[grain])
    ratio[grain] = cstar[grain]/mustar[grain]
    dratio[grain] = ratio[grain] * np.sqrt((dcstar[grain]/cstar[grain])**2 + (dmustar[grain]/mustar[grain])**2)
    Co[grain] = df_VES_2022.CD_VES[grain]
    CDsettle[grain] = 4*(df_density_2022.total_grain_ro[grain]/rof - 1)*g*df_shape_2022.dn[grain] / (3*df_vels_2022.vel_mean[grain]**2)
    CD[grain] = df_vels_2022.CD[grain]
    # cstar[grain] = ((df_VES_2022.ws_sphere[grain])/(df_vels_2022.vel_mean[grain]))**2 * df_shape_2022.CSF[grain]
    # mugb = (np.tan(np.radians(aor_gb)) - np.tan(np.radians(bed_angle)))
    # mustar[grain] = (np.tan(np.radians(df_ars_2022.mean_angle[grain])) - np.tan(np.radians(bed_angle)))/mugb
df_2 = pd.DataFrame({'Co': df_VES_2022.CD_VES, 'CD': df_vels_2022.CD, 'mustar': mustar, 'Cstar': cstar, 'dmustar': dmustar, 'dcstar': dcstar, 'ratio': ratio, 'dratio': dratio})

df = pd.concat([df_1, df_2])
# df['Cstar'] = df.CD/df.Co
# df['E'] = df.Cstar/df.mustar
# sns.distplot(df.E, kde=False)
# df
# %%
param_list = ['grain_density', 'a', 'b', 'c', 'do', 'ws', 'wo', 'CDsettle', 'CD', 'Co', 'Rep', 'AoR', 'mus', 'Sf', 'Cstar', 'mustar', 'ratio']
grains = ['gb', 'oc', 'gg', 'ng', 'ls']
for i, grain in enumerate(grains):
    p_out = np.empty([len(param_list)]).astype(object)
    param_values = {'grain_density': df_density.total_grain_ro[grain], 'a': df_shape.A[grain]*1000, 'b': df_shape.B[grain]*1000, 'c': df_shape.C[grain]*1000, 'do': df_shape.dn[grain]*1000, 'ws': df_vels.vel_mean[grain], 'wo': df_VES.ws_sphere[grain], 'CDsettle': df_vels.CDsettle[grain], 'CD': df_vels.CD[grain], 'Co': df_VES.CD_VES[grain], 'Rep': df_VES.Re[grain], 'AoR': df_ars.mean_angle[grain], 'mus': np.tan(np.radians(df_ars.mean_angle[grain])), 'Sf': df_shape.CSF[grain], 'Cstar': df.Cstar[grain], 'mustar': df.mustar[grain], 'ratio': df.ratio[grain]}
    error_values = {'grain_density': df_density.dtotal_grain_ro[grain], 'a': df_shape.dA[grain]*1000, 'b': df_shape.dB[grain]*1000, 'c': df_shape.dC[grain]*1000, 'do': df_shape.ddn[grain]*1000, 'ws': df_vels.dvel_mean[grain], 'wo': df_VES.dws_sphere[grain], 'CDsettle': df_vels.CDsettle[grain], 'CD': df_vels.CD[grain], 'Co': df_VES.CD_VES[grain], 'Rep': df_VES.Re[grain], 'AoR': df_ars.std_angle[grain], 'mus': np.tan(np.radians(df_ars.std_angle[grain])), 'Sf': df_shape.dCSF[grain], 'Cstar': df.dcstar[grain], 'mustar': df.dmustar[grain], 'ratio': df.dratio[grain]}
    value_string = {'grain_density': '%4.0f', 'a': '%1.1f', 'b': '%1.1f', 'c': '%1.1f', 'do': '%1.1f', 'ws': '%1.3f', 'wo': '%1.3f', 'CDsettle': '%1.2f', 'CD': '%1.2f', 'Co': '%1.2f', 'Rep': '%0.0f', 'AoR': '%2.1f', 'mus': '%1.2f', 'Sf': '%1.2f', 'Cstar': '%1.2f', 'mustar': '%1.2f', 'ratio': '%1.2f'}
    error_string = {'grain_density': '%3.0f', 'a': '%1.2f', 'b': '%1.2f', 'c': '%1.2f', 'do': '%1.2f', 'ws': '%1.4f', 'wo': '%1.4f', 'CDsettle': '%1.2f', 'CD': '%1.2f', 'Co': '%1.2f', 'Rep': '%0.0f', 'AoR': '%1.2f', 'mus': '%1.2f', 'Sf': '%1.2f', 'Cstar': '%1.2f', 'mustar': '%1.2f', 'ratio': '%1.2f'}
    for j, param in enumerate(param_list):
        if param not in ['CDsettle', 'CD', 'Co', 'Rep']:
            p_out[j] = value_string[param] % param_values[param] + ' ± ' + error_string[param] % error_values[param]
        else:
            p_out[j] = value_string[param] % param_values[param]

    if i == 0:
        df_out = pd.DataFrame({grain: p_out}, index=param_list).T
    else:
        df_out = pd.concat([df_out, pd.DataFrame({grain: p_out}, index=param_list).T])

grains = ['tg', 'sf', 'ng2', 'ng3', 'ng4', 'ng5']
for i, grain in enumerate(grains):
    p_out = np.empty([len(param_list)]).astype(object)
    param_values = {'grain_density': df_density_2022.total_grain_ro[grain], 'a': df_shape_2022.A[grain]*1000, 'b': df_shape_2022.B[grain]*1000, 'c': df_shape_2022.C[grain]*1000, 'do': df_shape_2022.dn[grain]*1000, 'ws': df_vels_2022.vel_mean[grain], 'wo': df_VES_2022.ws_sphere[grain], 'CDsettle': df_vels_2022.CDsettle[grain], 'CD': df_vels_2022.CD[grain], 'Co': df_VES_2022.CD_VES[grain], 'Rep': df_VES_2022.Re[grain], 'AoR': df_ars_2022.mean_angle[grain], 'mus': np.tan(np.radians(df_ars_2022.mean_angle[grain])), 'Sf': df_shape_2022.CSF[grain], 'Cstar': df.Cstar[grain], 'mustar': df.mustar[grain], 'ratio': df.ratio[grain]}
    error_values = {'grain_density': df_density_2022.dtotal_grain_ro[grain], 'a': df_shape_2022.dA[grain]*1000, 'b': df_shape_2022.dB[grain]*1000, 'c': df_shape_2022.dC[grain]*1000, 'do': df_shape_2022.ddn[grain]*1000, 'ws': df_vels_2022.dvel_mean[grain], 'wo': df_VES_2022.dws_sphere[grain], 'CDsettle': df_vels_2022.CDsettle[grain], 'CD': df_vels_2022.CD[grain], 'Co': df_VES_2022.CD_VES[grain], 'Rep': df_VES_2022.Re[grain], 'AoR': df_ars_2022.std_angle[grain], 'mus': np.tan(np.radians(df_ars_2022.std_angle[grain])), 'Sf': df_shape_2022.dCSF[grain], 'Cstar': df.dcstar[grain], 'mustar': df.dmustar[grain], 'ratio': df.dratio[grain]}
    value_string = {'grain_density': '%4.0f', 'a': '%1.1f', 'b': '%1.1f', 'c': '%1.1f', 'do': '%1.1f', 'ws': '%1.3f', 'wo': '%1.3f', 'CDsettle': '%1.2f', 'CD': '%1.2f', 'Co': '%1.2f', 'Rep': '%0.0f', 'AoR': '%2.1f', 'mus': '%1.2f', 'Sf': '%1.2f', 'Cstar': '%1.2f', 'mustar': '%1.2f', 'ratio': '%1.2f'}
    error_string = {'grain_density': '%3.0f', 'a': '%1.2f', 'b': '%1.2f', 'c': '%1.2f', 'do': '%1.2f', 'ws': '%1.4f', 'wo': '%1.4f', 'CDsettle': '%1.2f', 'CD': '%1.2f', 'Co': '%1.2f', 'Rep': '%0.0f', 'AoR': '%1.2f', 'mus': '%1.2f', 'Sf': '%1.2f', 'Cstar': '%1.2f', 'mustar': '%1.2f', 'ratio': '%1.2f'}
    for j, param in enumerate(param_list):
        if param not in ['CDsettle', 'CD', 'Co', 'Rep']:
            p_out[j] = value_string[param] % param_values[param] + ' ± ' + error_string[param] % error_values[param]
        else:
            p_out[j] = value_string[param] % param_values[param]

    if i == 0:
        df_out_2 = pd.DataFrame({grain: p_out}, index=param_list).T
    else:
        df_out_2 = pd.concat([df_out_2, pd.DataFrame({grain: p_out}, index=param_list).T])

df_out.T.to_csv('table_main.csv')
df_out_2.T.to_csv('table_aux.csv')
# %%
